#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级下载CLI工具
支持断点续传、多线程下载、速度监控、P2P种子下载等功能

版本: 1.0.0
作者: Ghengzhiyiya
许可证: MIT
"""

__version__ = "1.0.0"
__author__ = "Ghengzhiyiya"
__license__ = "MIT"

import os
import sys
import time
import threading
import argparse
import requests
import hashlib
import json
from pathlib import Path
from urllib.parse import urlparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, asdict
from typing import Optional, List, Dict, Any
import signal
import configparser

# 尝试多种可能的libtorrent导入方式
TORRENT_SUPPORT = False
try:
    import libtorrent as lt
    TORRENT_SUPPORT = True
except ImportError:
    try:
        # 有些系统上可能是python-libtorrent
        import python_libtorrent as lt
        TORRENT_SUPPORT = True
    except ImportError:
        try:
            # 或者是python3-libtorrent
            from python3_libtorrent import lt
            TORRENT_SUPPORT = True
        except ImportError:
            pass  # 静默处理，不显示警告

@dataclass
class DownloadConfig:
    """下载配置类"""
    url: str
    output_path: str
    threads: int = 4
    chunk_size: int = 1024 * 1024  # 1MB
    timeout: int = 30
    retry_times: int = 3
    min_speed_threshold: float = 1024  # 1KB/s
    speed_check_interval: int = 10  # 秒
    resume: bool = True

class SpeedMonitor:
    """下载速度监控器"""
    
    def __init__(self, threshold: float, check_interval: int):
        self.threshold = threshold
        self.check_interval = check_interval
        self.downloaded_bytes = 0
        self.last_check_time = time.time()
        self.last_downloaded = 0
        self.speed_history = []
        self.is_slow = False
        
    def update(self, bytes_downloaded: int):
        """更新下载字节数"""
        self.downloaded_bytes = bytes_downloaded
        
    def check_speed(self) -> bool:
        """检查下载速度是否过慢"""
        current_time = time.time()
        if current_time - self.last_check_time >= self.check_interval:
            time_diff = current_time - self.last_check_time
            bytes_diff = self.downloaded_bytes - self.last_downloaded
            current_speed = bytes_diff / time_diff
            
            self.speed_history.append(current_speed)
            if len(self.speed_history) > 5:  # 保持最近5次记录
                self.speed_history.pop(0)
                
            # 如果连续3次速度都低于阈值，认为速度过慢
            if len(self.speed_history) >= 3:
                avg_speed = sum(self.speed_history[-3:]) / 3
                self.is_slow = avg_speed < self.threshold
                
            self.last_check_time = current_time
            self.last_downloaded = self.downloaded_bytes
            
            return current_speed
        return 0
        
    def format_speed(self, speed: float) -> str:
        """格式化速度显示"""
        if speed >= 1024 * 1024:
            return f"{speed / (1024 * 1024):.2f} MB/s"
        elif speed >= 1024:
            return f"{speed / 1024:.2f} KB/s"
        else:
            return f"{speed:.2f} B/s"

class DownloadProgress:
    """下载进度管理器"""
    
    def __init__(self, total_size: int, filename: str):
        self.total_size = total_size
        self.filename = filename
        self.downloaded = 0
        self.start_time = time.time()
        self.lock = threading.Lock()
        
    def update(self, chunk_size: int):
        """更新下载进度"""
        with self.lock:
            self.downloaded += chunk_size
            
    def get_progress(self) -> Dict[str, Any]:
        """获取进度信息"""
        with self.lock:
            elapsed = time.time() - self.start_time
            speed = self.downloaded / elapsed if elapsed > 0 else 0
            percentage = (self.downloaded / self.total_size * 100) if self.total_size > 0 else 0
            
            return {
                'downloaded': self.downloaded,
                'total': self.total_size,
                'percentage': percentage,
                'speed': speed,
                'elapsed': elapsed,
                'filename': self.filename
            }
            
    def display(self):
        """显示进度条"""
        info = self.get_progress()
        bar_length = 50
        filled_length = int(bar_length * info['percentage'] / 100)
        bar = '█' * filled_length + '-' * (bar_length - filled_length)
        
        speed_str = SpeedMonitor(0, 0).format_speed(info['speed'])
        
        print(f"\r{info['filename']}: |{bar}| {info['percentage']:.1f}% "
              f"({self._format_size(info['downloaded'])}/{self._format_size(info['total'])}) "
              f"速度: {speed_str}", end='', flush=True)
              
    def _format_size(self, size: int) -> str:
        """格式化文件大小"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"

class MultiThreadDownloader:
    """多线程下载器"""
    
    def __init__(self, config: DownloadConfig):
        self.config = config
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
        self.progress = None
        self.speed_monitor = None
        self.is_paused = False
        self.is_cancelled = False
        
    def get_file_info(self) -> Dict[str, Any]:
        """获取文件信息"""
        try:
            response = self.session.head(self.config.url, timeout=self.config.timeout)
            response.raise_for_status()
            
            # 获取文件大小
            content_length = response.headers.get('Content-Length')
            total_size = int(content_length) if content_length else 0
            
            # 检查是否支持断点续传
            accept_ranges = response.headers.get('Accept-Ranges', '').lower()
            supports_resume = accept_ranges == 'bytes'
            
            # 获取文件名
            filename = self._get_filename_from_response(response)
            
            return {
                'total_size': total_size,
                'supports_resume': supports_resume,
                'filename': filename
            }
        except Exception as e:
            raise Exception(f"获取文件信息失败: {e}")
            
    def _get_filename_from_response(self, response) -> str:
        """从响应中获取文件名"""
        # 尝试从Content-Disposition获取
        content_disposition = response.headers.get('Content-Disposition', '')
        if 'filename=' in content_disposition:
            filename = content_disposition.split('filename=')[1].strip('"\'')
            return filename
            
        # 从URL获取
        parsed_url = urlparse(self.config.url)
        filename = os.path.basename(parsed_url.path)
        
        if not filename or '.' not in filename:
            filename = 'download_file'
            
        return filename
        
    def download_chunk(self, start: int, end: int, chunk_id: int, temp_file: str) -> bool:
        """下载文件块"""
        headers = {'Range': f'bytes={start}-{end}'}
        
        for attempt in range(self.config.retry_times):
            if self.is_cancelled:
                return False
                
            try:
                response = self.session.get(
                    self.config.url, 
                    headers=headers, 
                    stream=True, 
                    timeout=self.config.timeout
                )
                response.raise_for_status()
                
                with open(temp_file, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if self.is_cancelled:
                            return False
                            
                        while self.is_paused:
                            time.sleep(0.1)
                            
                        if chunk:
                            f.write(chunk)
                            if self.progress:
                                self.progress.update(len(chunk))
                            if self.speed_monitor:
                                self.speed_monitor.update(self.progress.downloaded)
                                
                return True
                
            except Exception as e:
                print(f"\n块 {chunk_id} 下载失败 (尝试 {attempt + 1}/{self.config.retry_times}): {e}")
                if attempt == self.config.retry_times - 1:
                    return False
                time.sleep(2 ** attempt)  # 指数退避
                
        return False
        
    def download(self) -> bool:
        """执行下载"""
        try:
            # 获取文件信息
            file_info = self.get_file_info()
            total_size = file_info['total_size']
            supports_resume = file_info['supports_resume']
            filename = file_info['filename']
            
            # 确定输出路径
            if os.path.isdir(self.config.output_path):
                output_file = os.path.join(self.config.output_path, filename)
            else:
                output_file = self.config.output_path
                
            print(f"开始下载: {filename}")
            print(f"文件大小: {self._format_size(total_size)}")
            print(f"支持断点续传: {'是' if supports_resume else '否'}")
            print(f"线程数: {self.config.threads}")
            
            # 检查是否已存在部分下载的文件
            resume_info = self._load_resume_info(output_file)
            start_pos = 0
            
            if self.config.resume and resume_info and supports_resume:
                if resume_info.get('total_size') == total_size:
                    start_pos = resume_info.get('downloaded', 0)
                    print(f"恢复下载，已下载: {self._format_size(start_pos)}")
                    
            # 初始化进度和速度监控
            self.progress = DownloadProgress(total_size, filename)
            self.progress.downloaded = start_pos
            
            self.speed_monitor = SpeedMonitor(
                self.config.min_speed_threshold,
                self.config.speed_check_interval
            )
            
            # 如果不支持多线程或文件较小，使用单线程下载
            if not supports_resume or total_size < self.config.chunk_size * 2:
                return self._single_thread_download(output_file, start_pos)
            else:
                return self._multi_thread_download(output_file, total_size, start_pos)
                
        except Exception as e:
            print(f"\n下载失败: {e}")
            return False
            
    def _single_thread_download(self, output_file: str, start_pos: int = 0) -> bool:
        """单线程下载"""
        headers = {}
        if start_pos > 0:
            headers['Range'] = f'bytes={start_pos}-'
            
        mode = 'ab' if start_pos > 0 else 'wb'
        
        try:
            response = self.session.get(
                self.config.url, 
                headers=headers, 
                stream=True, 
                timeout=self.config.timeout
            )
            response.raise_for_status()
            
            with open(output_file, mode) as f:
                for chunk in response.iter_content(chunk_size=self.config.chunk_size):
                    if self.is_cancelled:
                        return False
                        
                    while self.is_paused:
                        time.sleep(0.1)
                        
                    if chunk:
                        f.write(chunk)
                        self.progress.update(len(chunk))
                        self.speed_monitor.update(self.progress.downloaded)
                        
                        # 检查速度
                        current_speed = self.speed_monitor.check_speed()
                        if current_speed > 0:
                            self.progress.display()
                            
                        # 如果速度过慢，暂停下载
                        if self.speed_monitor.is_slow:
                            print(f"\n检测到下载速度过慢，暂停下载...")
                            self._save_resume_info(output_file)
                            return False
                            
            print(f"\n下载完成: {output_file}")
            self._cleanup_resume_info(output_file)
            return True
            
        except Exception as e:
            print(f"\n单线程下载失败: {e}")
            self._save_resume_info(output_file)
            return False
            
    def _multi_thread_download(self, output_file: str, total_size: int, start_pos: int = 0) -> bool:
        """多线程下载"""
        remaining_size = total_size - start_pos
        chunk_size = remaining_size // self.config.threads
        
        # 创建临时文件
        temp_dir = f"{output_file}.tmp"
        os.makedirs(temp_dir, exist_ok=True)
        
        # 计算每个线程的下载范围
        download_tasks = []
        for i in range(self.config.threads):
            start = start_pos + i * chunk_size
            end = start + chunk_size - 1
            if i == self.config.threads - 1:  # 最后一个线程下载剩余部分
                end = total_size - 1
                
            temp_file = os.path.join(temp_dir, f"chunk_{i}")
            download_tasks.append((start, end, i, temp_file))
            
        print(f"使用 {self.config.threads} 个线程下载...")
        
        # 启动多线程下载
        with ThreadPoolExecutor(max_workers=self.config.threads) as executor:
            futures = {
                executor.submit(self.download_chunk, start, end, chunk_id, temp_file): chunk_id
                for start, end, chunk_id, temp_file in download_tasks
            }
            
            # 监控进度
            progress_thread = threading.Thread(target=self._monitor_progress)
            progress_thread.daemon = True
            progress_thread.start()
            
            # 等待所有任务完成
            success_count = 0
            for future in as_completed(futures):
                if future.result():
                    success_count += 1
                    
        # 检查是否所有块都下载成功
        if success_count != self.config.threads:
            print(f"\n部分下载块失败，保存断点信息...")
            self._save_resume_info(output_file)
            return False
            
        # 合并文件块
        print(f"\n合并文件块...")
        try:
            with open(output_file, 'ab' if start_pos > 0 else 'wb') as output:
                if start_pos == 0:  # 如果是全新下载，清空文件
                    output.seek(0)
                    output.truncate()
                    
                for i in range(self.config.threads):
                    temp_file = os.path.join(temp_dir, f"chunk_{i}")
                    if os.path.exists(temp_file):
                        with open(temp_file, 'rb') as chunk_file:
                            output.write(chunk_file.read())
                        os.remove(temp_file)
                        
            # 清理临时目录
            os.rmdir(temp_dir)
            
            print(f"\n下载完成: {output_file}")
            self._cleanup_resume_info(output_file)
            return True
            
        except Exception as e:
            print(f"\n合并文件失败: {e}")
            return False
            
    def _monitor_progress(self):
        """监控下载进度"""
        while not self.is_cancelled and self.progress.downloaded < self.progress.total_size:
            self.progress.display()
            
            # 检查速度
            current_speed = self.speed_monitor.check_speed()
            
            # 如果速度过慢，设置暂停标志
            if self.speed_monitor.is_slow:
                print(f"\n检测到下载速度过慢，暂停下载...")
                self.is_paused = True
                break
                
            time.sleep(1)
            
    def _save_resume_info(self, output_file: str):
        """保存断点续传信息"""
        resume_file = f"{output_file}.resume"
        resume_info = {
            'url': self.config.url,
            'total_size': self.progress.total_size,
            'downloaded': self.progress.downloaded,
            'timestamp': time.time()
        }
        
        try:
            with open(resume_file, 'w', encoding='utf-8') as f:
                json.dump(resume_info, f)
        except Exception as e:
            print(f"保存断点信息失败: {e}")
            
    def _load_resume_info(self, output_file: str) -> Optional[Dict]:
        """加载断点续传信息"""
        resume_file = f"{output_file}.resume"
        
        if not os.path.exists(resume_file):
            return None
            
        try:
            with open(resume_file, 'r', encoding='utf-8') as f:
                resume_info = json.load(f)
                
            # 检查URL是否匹配
            if resume_info.get('url') != self.config.url:
                return None
                
            # 检查时间戳（超过24小时的断点信息可能无效）
            if time.time() - resume_info.get('timestamp', 0) > 24 * 3600:
                return None
                
            return resume_info
            
        except Exception:
            return None
            
    def _cleanup_resume_info(self, output_file: str):
        """清理断点续传信息"""
        resume_file = f"{output_file}.resume"
        if os.path.exists(resume_file):
            try:
                os.remove(resume_file)
            except Exception:
                pass
                
    def _format_size(self, size: int) -> str:
        """格式化文件大小"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"
        
    def pause(self):
        """暂停下载"""
        self.is_paused = True
        
    def resume(self):
        """恢复下载"""
        self.is_paused = False
        
    def cancel(self):
        """取消下载"""
        self.is_cancelled = True

class TorrentDownloader:
    """P2P种子下载器"""
    
    def __init__(self, config: DownloadConfig):
        if not TORRENT_SUPPORT:
            raise Exception("P2P种子下载需要安装libtorrent库")
            
        self.config = config
        self.session = None
        self.handle = None
        
    def download(self, torrent_file: str) -> bool:
        """下载种子文件"""
        try:
            # 创建libtorrent会话
            self.session = lt.session()
            self.session.listen_on(6881, 6891)
            
            # 添加种子
            if torrent_file.startswith('magnet:'):
                # 磁力链接
                self.handle = lt.add_magnet_uri(self.session, torrent_file, {
                    'save_path': self.config.output_path
                })
            else:
                # 种子文件
                info = lt.torrent_info(torrent_file)
                self.handle = self.session.add_torrent({
                    'ti': info,
                    'save_path': self.config.output_path
                })
                
            print(f"开始P2P下载: {self.handle.name()}")
            
            # 监控下载进度
            while not self.handle.is_seed():
                status = self.handle.status()
                
                print(f"\r进度: {status.progress * 100:.1f}% "
                      f"下载速度: {status.download_rate / 1024:.1f} KB/s "
                      f"上传速度: {status.upload_rate / 1024:.1f} KB/s "
                      f"连接数: {status.num_peers}", end='', flush=True)
                      
                time.sleep(1)
                
            print(f"\nP2P下载完成: {self.handle.name()}")
            return True
            
        except Exception as e:
            print(f"\nP2P下载失败: {e}")
            return False
        finally:
            if self.session:
                self.session.pause()

def load_config_file(config_path: str) -> Dict[str, Any]:
    """加载配置文件"""
    config = configparser.ConfigParser()
    defaults = {
        'threads': '4',
        'chunk_size': '1048576',
        'timeout': '30',
        'retry': '3',
        'min_speed': '1024',
        'speed_check': '10',
        'resume': 'true'
    }
    
    if os.path.exists(config_path):
        try:
            config.read(config_path, encoding='utf-8')
            if 'download' in config:
                return dict(config['download'])
        except Exception as e:
            print(f"警告: 读取配置文件失败: {e}")
    
    return defaults

def save_config_file(config_path: str, config_dict: Dict[str, Any]):
    """保存配置文件"""
    config = configparser.ConfigParser()
    config['download'] = {k: str(v) for k, v in config_dict.items()}
    
    try:
        with open(config_path, 'w', encoding='utf-8') as f:
            config.write(f)
        print(f"配置已保存到: {config_path}")
    except Exception as e:
        print(f"警告: 保存配置文件失败: {e}")

def signal_handler(signum, frame):
    """信号处理器"""
    print("\n接收到中断信号，正在清理...")
    sys.exit(0)

def find_default_config() -> Optional[str]:
    """查找默认配置文件"""
    # 按优先级查找配置文件
    config_paths = [
        'download.conf',  # 当前目录
        os.path.expanduser('~/.download.conf'),  # 用户主目录
        os.path.expanduser('~/download.conf'),   # 用户主目录（备选）
    ]
    
    for path in config_paths:
        if os.path.exists(path):
            return path
    return None

def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='高级下载CLI工具 - 自动加载配置文件',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=f"""
示例用法:
  %(prog)s https://example.com/file.zip -o ~/Downloads/
  %(prog)s https://example.com/file.zip -t 8 -c 2097152
  %(prog)s file.torrent --torrent -o /download/path/
  %(prog)s "magnet:?xt=urn:btih:..." -o /download/path/

配置文件:
  程序会自动查找并加载以下配置文件（按优先级）:
  1. ./download.conf (当前目录)
  2. ~/.download.conf (用户主目录)
  3. ~/download.conf (用户主目录备选)
  
  使用 --config 可指定自定义配置文件
  命令行参数优先级高于配置文件设置

版本: {__version__}
作者: {__author__}
许可证: {__license__}
        """
    )
    
    parser.add_argument('url', help='下载URL或种子文件路径')
    parser.add_argument('-o', '--output', default='.', help='输出路径')
    parser.add_argument('-t', '--threads', type=int, help='线程数 (1-32)')
    parser.add_argument('-c', '--chunk-size', type=int, help='块大小(字节)')
    parser.add_argument('--timeout', type=int, help='超时时间(秒)')
    parser.add_argument('--retry', type=int, help='重试次数')
    parser.add_argument('--min-speed', type=float, help='最小速度阈值(字节/秒)')
    parser.add_argument('--speed-check', type=int, help='速度检查间隔(秒)')
    parser.add_argument('--no-resume', action='store_true', help='禁用断点续传')
    parser.add_argument('--torrent', action='store_true', help='P2P种子下载模式')
    parser.add_argument('--config', help='指定配置文件路径（覆盖自动查找）')
    parser.add_argument('--save-config', help='保存当前配置到指定文件')
    parser.add_argument('--no-config', action='store_true', help='不加载任何配置文件')
    parser.add_argument('--version', action='version', version=f'%(prog)s {__version__}')
    
    args = parser.parse_args()
    
    # 自动查找并加载配置文件
    config_defaults = {}
    config_file_used = None
    
    if not args.no_config:
        if args.config:
            # 用户指定的配置文件
            config_file_used = args.config
        else:
            # 自动查找默认配置文件
            config_file_used = find_default_config()
        
        if config_file_used:
            config_defaults = load_config_file(config_file_used)
            print(f"已加载配置文件: {config_file_used}")
        else:
            print("未找到配置文件，使用默认设置")
    
    # 应用配置文件默认值（命令行参数优先）
    args.threads = args.threads if args.threads is not None else int(config_defaults.get('threads', 4))
    args.chunk_size = args.chunk_size if args.chunk_size is not None else int(config_defaults.get('chunk_size', 1024*1024))
    args.timeout = args.timeout if args.timeout is not None else int(config_defaults.get('timeout', 30))
    args.retry = args.retry if args.retry is not None else int(config_defaults.get('retry', 3))
    args.min_speed = args.min_speed if args.min_speed is not None else float(config_defaults.get('min_speed', 1024))
    args.speed_check = args.speed_check if args.speed_check is not None else int(config_defaults.get('speed_check', 10))
    
    # 处理resume参数（--no-resume优先级最高）
    if not args.no_resume:
        args.no_resume = not config_defaults.get('resume', 'true').lower() == 'true'
    
    # 参数验证
    if args.threads < 1 or args.threads > 32:
        print("错误: 线程数必须在1-32之间")
        sys.exit(1)
        
    if args.chunk_size < 1024:
        print("错误: 块大小不能小于1024字节")
        sys.exit(1)
        
    if args.timeout < 1:
        print("错误: 超时时间必须大于0秒")
        sys.exit(1)
        
    if args.retry < 1:
        print("错误: 重试次数必须大于0")
        sys.exit(1)
        
    if args.min_speed < 0:
        print("错误: 最小速度阈值不能为负数")
        sys.exit(1)
        
    if args.speed_check < 1:
        print("错误: 速度检查间隔必须大于0秒")
        sys.exit(1)
    
    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    # 创建下载配置
    config = DownloadConfig(
        url=args.url,
        output_path=args.output,
        threads=args.threads,
        chunk_size=args.chunk_size,
        timeout=args.timeout,
        retry_times=args.retry,
        min_speed_threshold=args.min_speed,
        speed_check_interval=args.speed_check,
        resume=not args.no_resume
    )
    
    try:
        if args.torrent or args.url.endswith('.torrent') or args.url.startswith('magnet:'):
            # P2P种子下载
            downloader = TorrentDownloader(config)
            success = downloader.download(args.url)
        else:
            # HTTP下载
            downloader = MultiThreadDownloader(config)
            success = downloader.download()
            
        # 保存配置（如果指定）
        if args.save_config:
            config_to_save = {
                'threads': args.threads,
                'chunk_size': args.chunk_size,
                'timeout': args.timeout,
                'retry': args.retry,
                'min_speed': args.min_speed,
                'speed_check': args.speed_check,
                'resume': not args.no_resume
            }
            save_config_file(args.save_config, config_to_save)
        
        if success:
            print("\n下载成功！")
            sys.exit(0)
        else:
            print("\n下载失败！")
            sys.exit(1)
            
    except KeyboardInterrupt:
        print("\n用户取消下载")
        sys.exit(1)
    except Exception as e:
        print(f"\n错误: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()